#pragma once
#include <Eigen/Core>
#include <ceres/ceres.h>
#include <ceres/autodiff_cost_function.h>
#include <ceres/cubic_interpolation.h>
#include <ceres/rotation.h>

class PoseGraph3dErrorTerm {
 public:
  PoseGraph3dErrorTerm(const Eigen::Vector2d& point2d,const double fx,
                       const double fy, const double cx, const double cy,
                       const double obs_num)
      : point2d_(point2d),
        fx_(fx),
        fy_(fy),
        cx_(cx),
        cy_(cy),
        weight_(obs_num) {}

  template <typename T>
  bool operator()(const T* const cam_q, const T* const cam_t,
                  const T* const land_mark_3d, T* residuals_ptr) const {
  
    const Eigen::Quaternion<T> q(cam_q[3], cam_q[0], cam_q[1], cam_q[2]);
    Eigen::Map<const Eigen::Matrix<T, 3, 1>> t(cam_t);
    Eigen::Map<const Eigen::Matrix<T, 3, 1>> point3d(land_mark_3d);

    Eigen::Matrix<T, 3, 1> p_cp =
        q.conjugate().toRotationMatrix() * (point3d - t);
    T u = T(fx_) * p_cp(0) / p_cp(2) + T(cx_);
    T v = T(fy_) * p_cp(1) / p_cp(2) + T(cy_);

     // Compute the map point pose in pixel frame.
    residuals_ptr[0] = (u - T(point2d_[0])) * T(weight_);
    residuals_ptr[1] = (v - T(point2d_[1])) * T(weight_);
    return true;
  }

  static ceres::CostFunction* Create(const Eigen::Vector2d& point2d,
                                     const Eigen::Matrix3d& K,
                                     const double& sqrt_information) {
    double fx = K(0, 0);
    double fy = K(1, 1);
    double cx = K(0, 2);
    double cy = K(1, 2);
    return new ceres::AutoDiffCostFunction<PoseGraph3dErrorTerm, 2, 4, 3, 3>(
        new PoseGraph3dErrorTerm(point2d, fx, fy, cx, cy, sqrt_information));
  }
  EIGEN_MAKE_ALIGNED_OPERATOR_NEW

 private:
  // The measurement for the position of B relative to A in the A frame.
  const Eigen::Vector2d point2d_;
  // The square root of the measurement information matrix.
  const double weight_;
  //
  const double fx_;
  //
  const double fy_;
  //
  const double cx_;
  //
  const double cy_;
};


class GlobalPriorPoseFactor {
 public:
  GlobalPriorPoseFactor(const Eigen::Vector3d& prior_t_,
                        const Eigen::Quaterniond& prior_q_,
                        const std::vector<double>& var_vec_)
      : prior_q(prior_q_), prior_t(prior_t_), var_vec(var_vec_) {}

  template <typename T>
  bool operator()(const T* const cam_quat, const T* const cam_tran,
                  T* residue) const {
    // compute residuals
    Eigen::Matrix<T, 3, 1> twc(cam_tran[0], cam_tran[1], cam_tran[2]);
    Eigen::Quaternion<T> qwc(cam_quat[3], cam_quat[0], cam_quat[1],
                             cam_quat[2]);

    Eigen::Quaternion<T> q_error = qwc.inverse() * prior_q.template cast<T>();
    T q_err_coeffs[4] = {q_error.w(), q_error.x(), q_error.y(), q_error.z()};
    Eigen::Matrix<T, 3, 3> R_err = QuaternionToRotation<T>(q_err_coeffs);
    Eigen::Matrix<T, 3, 1> t_error = twc - prior_t.template cast<T>();

    T roll, pitch, yaw;
    mat2RPY(R_err, roll, pitch, yaw);

    Eigen::Map<Eigen::Matrix<T, 6, 1>> residuals(residue);
    residuals[0] = t_error(0) / T(var_vec[0]);
    residuals[1] = t_error(1) / T(var_vec[1]);
    residuals[2] = t_error(2) / T(var_vec[2]);
    residuals[3] = roll / T(var_vec[3]);
    residuals[4] = pitch / T(var_vec[4]);
    residuals[5] = yaw / T(var_vec[5]);
    return true;
  }

  template <typename T>
  void mat2RPY(const Eigen::Matrix<T, 3, 3>& m, T& roll, T& pitch,
               T& yaw) const {
    roll = atan2(m(2, 1), m(2, 2));
    pitch = atan2(-m(2, 0), sqrt(m(2, 1) * m(2, 1) + m(2, 2) * m(2, 2)));
    yaw = atan2(m(1, 0), m(0, 0));
  }

  template <typename T>
  Eigen::Matrix<T, 3, 3> QuaternionToRotation(const T* const q) const {
    T R[9];
    ceres::QuaternionToRotation(q, R);

    Eigen::Matrix<T, 3, 3> rmat;
    for (int i = 0; i < 3; ++i) {
      for (int j = 0; j < 3; ++j) {
        rmat(i, j) = R[i * 3 + j];
      }
    }
    return rmat;
  }

  static ceres::CostFunction* Create(const Eigen::Matrix4d& prior_p,
                                     const std::vector<double>& var_vec_) {
    Eigen::Vector3d prior_t = prior_p.block(0, 3, 3, 1);
    Eigen::Quaterniond prior_q(prior_p.block<3, 3>(0, 0));
    return (
        //  error dimention 1
        new ceres::AutoDiffCostFunction<GlobalPriorPoseFactor, 6, 4, 3>(
            new GlobalPriorPoseFactor(prior_t, prior_q, var_vec_)));
  }

 private:
  const Eigen::Vector3d prior_t;
  const Eigen::Quaterniond prior_q;
  const std::vector<double> var_vec;
};

class PriorDeltaErrorTerm {
 public:
  PriorDeltaErrorTerm(const Eigen::Vector3d& delta_t_,
                      const Eigen::Quaterniond& delta_q_,
                      double trans_information_, double rot_information_)
      : delta_q(delta_q_),
        delta_t(delta_t_),
        trans_information(trans_information_),
        rot_information(rot_information_) {}

  template <typename T>
  bool operator()(const T* const quat_i, const T* const tran_i,
                  const T* const quat_j, const T* const tran_j,
                  T* residuals_ptr) const {
    const Eigen::Matrix<T, 3, 1> twc_i(tran_i[0], tran_i[1], tran_i[2]);
    const Eigen::Matrix<T, 3, 1> twc_j(tran_j[0], tran_j[1], tran_j[2]);
    Eigen::Quaternion<T> Qwc_i(quat_i[3], quat_i[0], quat_i[1], quat_i[2]);
    Eigen::Quaternion<T> Qwc_j(quat_j[3], quat_j[0], quat_j[1], quat_j[2]);

    Eigen::Quaternion<T> q_i_j = Qwc_i.conjugate() * Qwc_j;
    Eigen::Matrix<T, 3, 1> t_i_j = Qwc_i.conjugate() * (twc_j - twc_i);

    Eigen::Matrix<T, 3, 1> t_error = delta_t.template cast<T>() - t_i_j;
    Eigen::Quaternion<T> q_error =
        q_i_j.conjugate() * delta_q.template cast<T>();

    Eigen::Map<Eigen::Matrix<T, 6, 1>> residuals(residuals_ptr);
    residuals[0] = t_error[0] / T(trans_information);
    residuals[1] = t_error[1] / T(trans_information);
    residuals[2] = t_error[2] / T(trans_information);
    residuals.template block<3, 1>(3, 0) =
        T(2.0) * q_error.vec() / rot_information;
    return true;
  }

  static ceres::CostFunction* Create(Eigen::Matrix4d& delta_p,
                                     double trans_information,
                                     double rot_information) {
    Eigen::Vector3d delta_t = delta_p.block(0, 3, 3, 1);
    Eigen::Quaterniond delta_q(delta_p.block<3, 3>(0, 0));
    return (new ceres::AutoDiffCostFunction<PriorDeltaErrorTerm, 6, 4, 3, 4, 3>(
        new PriorDeltaErrorTerm(delta_t, delta_q, trans_information,
                                rot_information)));
  }

 private:
  // The 6dof pose of B relative to A in the A frame.
  const Eigen::Vector3d delta_t;
  const Eigen::Quaterniond delta_q;
  // The inverse square root of the measurement covariance matrix.
  double trans_information;
  double rot_information;
};
